"""BLE Authentication Manager.

Handles PIN-based authentication, device pairing, and user access control
for secure BLE communications.
"""

import asyncio
import hashlib
import hmac
import json
import logging
import secrets
import time
from typing import Dict, List, Optional, Set, Callable, Any
from enum import Enum
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


# Authentication Configuration
AUTH_CONFIG = {
    "pin_length": 6,             # PIN code length
    "pin_attempts": 3,           # Maximum attempt count
    "lockout_time": 300,         # Lockout time (5 minutes)
    "session_timeout": 1800,     # Session timeout (30 minutes)
    "challenge_timeout": 300,    # Challenge timeout (5 minutes)
    "cleanup_interval": 60,      # Cleanup interval (1 minute)
    "hash_rounds": 10000,        # PBKDF2 rounds
    "salt_size": 16              # Salt size in bytes
}


class AuthState(Enum):
    """Authentication states."""
    UNAUTHENTICATED = 0    # Not authenticated
    PIN_CHALLENGE = 1      # PIN challenge in progress
    AUTHENTICATED = 2      # Authenticated
    PAIRED = 3            # Device paired
    EXPIRED = 4           # Session expired
    LOCKED = 5            # Account locked
    ERROR = 6             # Error state


class AuthResult(Enum):
    """Authentication results."""
    SUCCESS = 0           # Authentication successful
    INVALID_PIN = 1       # Invalid PIN code
    CHALLENGE_EXPIRED = 2 # Challenge expired
    TOO_MANY_ATTEMPTS = 3 # Too many failed attempts
    DEVICE_LOCKED = 4     # Device is locked
    SYSTEM_ERROR = 5      # System error


@dataclass
class AuthChallenge:
    """Authentication challenge data."""
    challenge_id: str
    device_id: str
    pin_hash: str
    salt: bytes
    created_time: float
    expires_at: float
    attempts: int = 0
    max_attempts: int = 3
    
    @property
    def is_expired(self) -> bool:
        """Check if challenge is expired."""
        return time.time() > self.expires_at
    
    @property
    def attempts_left(self) -> int:
        """Get remaining attempts."""
        return max(0, self.max_attempts - self.attempts)
    
    def increment_attempts(self) -> bool:
        """Increment attempt counter."""
        self.attempts += 1
        return self.attempts < self.max_attempts


@dataclass
class AuthSession:
    """Authentication session data."""
    session_id: str
    device_id: str
    state: AuthState
    created_time: float
    last_activity: float
    expires_at: float
    permissions: Set[str] = field(default_factory=set)
    data: Dict[str, Any] = field(default_factory=dict)
    
    @property
    def is_expired(self) -> bool:
        """Check if session is expired."""
        return time.time() > self.expires_at
    
    @property
    def idle_time(self) -> float:
        """Get idle time in seconds."""
        return time.time() - self.last_activity
    
    def touch(self):
        """Update last activity time."""
        self.last_activity = time.time()
        self.expires_at = time.time() + AUTH_CONFIG["session_timeout"]


@dataclass
class DeviceInfo:
    """Registered device information."""
    device_id: str
    device_name: str
    mac_address: str
    paired_time: float
    last_seen: float
    trust_level: int = 1  # 1=basic, 2=trusted, 3=admin
    is_blocked: bool = False
    failed_attempts: int = 0
    locked_until: float = 0.0
    
    @property
    def is_locked(self) -> bool:
        """Check if device is currently locked."""
        return time.time() < self.locked_until
    
    def lock_device(self, duration: int = 300):
        """Lock device for specified duration."""
        self.locked_until = time.time() + duration
        self.is_blocked = True
    
    def unlock_device(self):
        """Unlock device."""
        self.locked_until = 0.0
        self.is_blocked = False
        self.failed_attempts = 0


class DeviceRegistry:
    """Device registry for pairing and trust management."""
    
    def __init__(self):
        self.paired_devices: Dict[str, DeviceInfo] = {}
        self.trusted_devices: Set[str] = set()
        self.blocked_devices: Set[str] = set()
        
    def register_device(self, device_id: str, device_name: str, 
                       mac_address: str, trust_level: int = 1) -> DeviceInfo:
        """Register a new device."""
        device_info = DeviceInfo(
            device_id=device_id,
            device_name=device_name,
            mac_address=mac_address,
            paired_time=time.time(),
            last_seen=time.time(),
            trust_level=trust_level
        )
        
        self.paired_devices[device_id] = device_info
        
        if trust_level >= 2:
            self.trusted_devices.add(device_id)
            
        logger.info(f"Device registered: {device_id} ({device_name})")
        return device_info
    
    def get_device(self, device_id: str) -> Optional[DeviceInfo]:
        """Get device information."""
        return self.paired_devices.get(device_id)
    
    def is_paired(self, device_id: str) -> bool:
        """Check if device is paired."""
        return device_id in self.paired_devices
    
    def is_trusted(self, device_id: str) -> bool:
        """Check if device is trusted."""
        return device_id in self.trusted_devices
    
    def is_blocked(self, device_id: str) -> bool:
        """Check if device is blocked."""
        device = self.paired_devices.get(device_id)
        return device and device.is_blocked or device_id in self.blocked_devices
    
    def block_device(self, device_id: str, reason: str = "Security violation"):
        """Block a device."""
        self.blocked_devices.add(device_id)
        device = self.paired_devices.get(device_id)
        if device:
            device.is_blocked = True
            device.lock_device()
        logger.warning(f"Device blocked: {device_id}, reason: {reason}")
    
    def unblock_device(self, device_id: str):
        """Unblock a device."""
        self.blocked_devices.discard(device_id)
        device = self.paired_devices.get(device_id)
        if device:
            device.unlock_device()
        logger.info(f"Device unblocked: {device_id}")
    
    def revoke_device(self, device_id: str):
        """Revoke device pairing."""
        self.paired_devices.pop(device_id, None)
        self.trusted_devices.discard(device_id)
        self.blocked_devices.discard(device_id)
        logger.info(f"Device pairing revoked: {device_id}")
    
    def update_device_activity(self, device_id: str):
        """Update device last seen time."""
        device = self.paired_devices.get(device_id)
        if device:
            device.last_seen = time.time()
    
    def cleanup_expired_devices(self, max_age_days: int = 30):
        """Clean up devices not seen for specified days."""
        cutoff_time = time.time() - (max_age_days * 24 * 3600)
        expired_devices = []
        
        for device_id, device in self.paired_devices.items():
            if device.last_seen < cutoff_time:
                expired_devices.append(device_id)
        
        for device_id in expired_devices:
            self.revoke_device(device_id)
        
        return len(expired_devices)
    
    def get_device_statistics(self) -> Dict[str, Any]:
        """Get device registry statistics."""
        return {
            "total_devices": len(self.paired_devices),
            "trusted_devices": len(self.trusted_devices),
            "blocked_devices": len(self.blocked_devices),
            "active_devices": sum(1 for d in self.paired_devices.values() 
                                if time.time() - d.last_seen < 3600)
        }


class AuthManager:
    """BLE Authentication Manager."""
    
    def __init__(self):
        self.device_registry = DeviceRegistry()
        self.active_challenges: Dict[str, AuthChallenge] = {}
        self.active_sessions: Dict[str, AuthSession] = {}
        self.device_sessions: Dict[str, str] = {}  # device_id -> session_id
        
        # Event callbacks
        self.on_auth_success: Optional[Callable[[str, AuthSession], None]] = None
        self.on_auth_failure: Optional[Callable[[str, AuthResult], None]] = None
        self.on_device_paired: Optional[Callable[[str, DeviceInfo], None]] = None
        self.on_session_expired: Optional[Callable[[str], None]] = None
        
        # Background tasks
        self.cleanup_task: Optional[asyncio.Task] = None
        self.is_running = False
        
        logger.info("AuthManager initialized")
    
    async def start(self):
        """Start authentication manager."""
        if self.is_running:
            return
            
        self.is_running = True
        self.cleanup_task = asyncio.create_task(self._cleanup_loop())
        
        logger.info("AuthManager started")
    
    async def stop(self):
        """Stop authentication manager."""
        self.is_running = False
        
        if self.cleanup_task:
            self.cleanup_task.cancel()
            try:
                await self.cleanup_task
            except asyncio.CancelledError:
                pass
        
        logger.info("AuthManager stopped")
    
    def generate_pin_code(self) -> str:
        """Generate a random PIN code."""
        pin_length = AUTH_CONFIG["pin_length"]
        return ''.join([str(secrets.randbelow(10)) for _ in range(pin_length)])
    
    def hash_pin_code(self, pin_code: str, salt: bytes = None) -> tuple[str, bytes]:
        """Hash PIN code with salt."""
        if salt is None:
            salt = secrets.token_bytes(AUTH_CONFIG["salt_size"])
        
        pin_bytes = pin_code.encode('utf-8')
        hash_rounds = AUTH_CONFIG["hash_rounds"]
        
        # Use PBKDF2 for secure PIN hashing
        hash_obj = hashlib.pbkdf2_hmac('sha256', pin_bytes, salt, hash_rounds)
        pin_hash = hash_obj.hex()
        
        return pin_hash, salt
    
    def verify_pin_code(self, pin_code: str, pin_hash: str, salt: bytes) -> bool:
        """Verify PIN code against hash."""
        try:
            computed_hash, _ = self.hash_pin_code(pin_code, salt)
            return hmac.compare_digest(pin_hash, computed_hash)
        except Exception as e:
            logger.error(f"PIN verification error: {e}")
            return False
    
    async def create_auth_challenge(self, device_id: str, 
                                  custom_pin: str = None) -> AuthChallenge:
        """Create authentication challenge for device."""
        # Check if device is blocked
        if self.device_registry.is_blocked(device_id):
            raise ValueError(f"Device {device_id} is blocked")
        
        # Generate PIN and challenge
        pin_code = custom_pin or self.generate_pin_code()
        pin_hash, salt = self.hash_pin_code(pin_code)
        
        challenge_id = secrets.token_urlsafe(16)
        challenge = AuthChallenge(
            challenge_id=challenge_id,
            device_id=device_id,
            pin_hash=pin_hash,
            salt=salt,
            created_time=time.time(),
            expires_at=time.time() + AUTH_CONFIG["challenge_timeout"],
            max_attempts=AUTH_CONFIG["pin_attempts"]
        )
        
        # Remove any existing challenges for this device
        existing_challenges = [cid for cid, ch in self.active_challenges.items() 
                             if ch.device_id == device_id]
        for cid in existing_challenges:
            del self.active_challenges[cid]
        
        self.active_challenges[challenge_id] = challenge
        
        logger.info(f"Auth challenge created for device {device_id}: {challenge_id}")
        logger.debug(f"PIN code for challenge {challenge_id}: {pin_code}")
        
        return challenge
    
    async def verify_auth_response(self, challenge_id: str, 
                                 pin_response: str) -> tuple[AuthResult, Optional[AuthSession]]:
        """Verify authentication response."""
        challenge = self.active_challenges.get(challenge_id)
        if not challenge:
            return AuthResult.CHALLENGE_EXPIRED, None
        
        # Check if challenge is expired
        if challenge.is_expired:
            del self.active_challenges[challenge_id]
            return AuthResult.CHALLENGE_EXPIRED, None
        
        # Check device lock status
        device = self.device_registry.get_device(challenge.device_id)
        if device and device.is_locked:
            return AuthResult.DEVICE_LOCKED, None
        
        # Verify PIN
        if not self.verify_pin_code(pin_response, challenge.pin_hash, challenge.salt):
            challenge.increment_attempts()
            
            # Update device failed attempts
            if device:
                device.failed_attempts += 1
                if device.failed_attempts >= AUTH_CONFIG["pin_attempts"]:
                    device.lock_device(AUTH_CONFIG["lockout_time"])
                    logger.warning(f"Device locked due to failed attempts: {challenge.device_id}")
            
            if challenge.attempts >= challenge.max_attempts:
                del self.active_challenges[challenge_id]
                if self.on_auth_failure:
                    self.on_auth_failure(challenge.device_id, AuthResult.TOO_MANY_ATTEMPTS)
                return AuthResult.TOO_MANY_ATTEMPTS, None
            
            if self.on_auth_failure:
                self.on_auth_failure(challenge.device_id, AuthResult.INVALID_PIN)
            return AuthResult.INVALID_PIN, None
        
        # PIN is correct - create authenticated session
        session = await self._create_auth_session(challenge.device_id)
        
        # Clean up challenge
        del self.active_challenges[challenge_id]
        
        # Update device info
        if device:
            device.failed_attempts = 0
            device.last_seen = time.time()
        
        # Trigger success callback
        if self.on_auth_success:
            self.on_auth_success(challenge.device_id, session)
        
        logger.info(f"Authentication successful for device: {challenge.device_id}")
        return AuthResult.SUCCESS, session
    
    async def _create_auth_session(self, device_id: str) -> AuthSession:
        """Create authenticated session."""
        session_id = secrets.token_urlsafe(32)
        
        # Determine initial state and permissions
        if self.device_registry.is_paired(device_id):
            state = AuthState.PAIRED
            permissions = {"read", "write", "control"}
            if self.device_registry.is_trusted(device_id):
                permissions.add("admin")
        else:
            state = AuthState.AUTHENTICATED
            permissions = {"read"}
        
        session = AuthSession(
            session_id=session_id,
            device_id=device_id,
            state=state,
            created_time=time.time(),
            last_activity=time.time(),
            expires_at=time.time() + AUTH_CONFIG["session_timeout"],
            permissions=permissions
        )
        
        # Remove any existing session for this device
        existing_session_id = self.device_sessions.get(device_id)
        if existing_session_id and existing_session_id in self.active_sessions:
            del self.active_sessions[existing_session_id]
        
        # Store new session
        self.active_sessions[session_id] = session
        self.device_sessions[device_id] = session_id
        
        return session
    
    async def pair_device(self, device_id: str, device_name: str, 
                         mac_address: str, trust_level: int = 1) -> DeviceInfo:
        """Pair and register a device."""
        # Check if device has an active authenticated session
        session_id = self.device_sessions.get(device_id)
        if not session_id or session_id not in self.active_sessions:
            raise ValueError("Device must be authenticated before pairing")
        
        session = self.active_sessions[session_id]
        if session.state not in [AuthState.AUTHENTICATED, AuthState.PAIRED]:
            raise ValueError("Invalid session state for pairing")
        
        # Register device
        device_info = self.device_registry.register_device(
            device_id, device_name, mac_address, trust_level
        )
        
        # Update session state
        session.state = AuthState.PAIRED
        session.permissions.update({"write", "control"})
        if trust_level >= 2:
            session.permissions.add("admin")
        
        # Trigger pairing callback
        if self.on_device_paired:
            self.on_device_paired(device_id, device_info)
        
        logger.info(f"Device paired successfully: {device_id} ({device_name})")
        return device_info
    
    def get_session(self, session_id: str) -> Optional[AuthSession]:
        """Get session by ID."""
        session = self.active_sessions.get(session_id)
        if session and not session.is_expired:
            session.touch()
            return session
        return None
    
    def get_session_by_device(self, device_id: str) -> Optional[AuthSession]:
        """Get session by device ID."""
        session_id = self.device_sessions.get(device_id)
        if session_id:
            return self.get_session(session_id)
        return None
    
    def has_permission(self, session_id: str, permission: str) -> bool:
        """Check if session has specific permission."""
        session = self.get_session(session_id)
        return session and permission in session.permissions
    
    async def revoke_session(self, session_id: str) -> bool:
        """Revoke authentication session."""
        session = self.active_sessions.get(session_id)
        if not session:
            return False
        
        # Remove from both mappings
        del self.active_sessions[session_id]
        if session.device_id in self.device_sessions:
            del self.device_sessions[session.device_id]
        
        logger.info(f"Session revoked: {session_id} (device: {session.device_id})")
        return True
    
    async def revoke_device_access(self, device_id: str) -> bool:
        """Revoke all access for a device."""
        # Remove active session
        session_id = self.device_sessions.get(device_id)
        if session_id:
            await self.revoke_session(session_id)
        
        # Remove device pairing
        self.device_registry.revoke_device(device_id)
        
        # Remove any pending challenges
        pending_challenges = [cid for cid, ch in self.active_challenges.items() 
                             if ch.device_id == device_id]
        for cid in pending_challenges:
            del self.active_challenges[cid]
        
        logger.info(f"Device access revoked: {device_id}")
        return True
    
    async def _cleanup_loop(self):
        """Background cleanup of expired challenges and sessions."""
        while self.is_running:
            try:
                await asyncio.sleep(AUTH_CONFIG["cleanup_interval"])
                
                current_time = time.time()
                
                # Clean up expired challenges
                expired_challenges = [cid for cid, ch in self.active_challenges.items() 
                                    if ch.is_expired]
                for cid in expired_challenges:
                    del self.active_challenges[cid]
                
                # Clean up expired sessions
                expired_sessions = [sid for sid, sess in self.active_sessions.items() 
                                  if sess.is_expired]
                for sid in expired_sessions:
                    session = self.active_sessions[sid]
                    del self.active_sessions[sid]
                    if session.device_id in self.device_sessions:
                        del self.device_sessions[session.device_id]
                    
                    if self.on_session_expired:
                        self.on_session_expired(session.device_id)
                
                # Clean up old devices
                self.device_registry.cleanup_expired_devices()
                
                if expired_challenges or expired_sessions:
                    logger.debug(f"Cleanup: removed {len(expired_challenges)} challenges, "
                               f"{len(expired_sessions)} sessions")
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Cleanup loop error: {e}")
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get authentication statistics."""
        return {
            "auth_manager": {
                "is_running": self.is_running,
                "active_challenges": len(self.active_challenges),
                "active_sessions": len(self.active_sessions),
                "authenticated_devices": sum(1 for s in self.active_sessions.values() 
                                           if s.state in [AuthState.AUTHENTICATED, AuthState.PAIRED])
            },
            "device_registry": self.device_registry.get_device_statistics()
        }
    
    async def cleanup(self):
        """Clean up authentication manager."""
        await self.stop()
        self.active_challenges.clear()
        self.active_sessions.clear()
        self.device_sessions.clear()
        logger.info("AuthManager cleanup completed")


# Example usage and testing
if __name__ == "__main__":
    import asyncio
    
    async def test_auth_manager():
        """Test authentication manager functionality."""
        auth_manager = AuthManager()
        
        # Set up event callbacks
        def on_auth_success(device_id, session):
            print(f"✓ Authentication successful: {device_id} -> {session.session_id}")
        
        def on_auth_failure(device_id, result):
            print(f"✗ Authentication failed: {device_id} -> {result.name}")
        
        def on_device_paired(device_id, device_info):
            print(f"🔗 Device paired: {device_id} -> {device_info.device_name}")
        
        auth_manager.on_auth_success = on_auth_success
        auth_manager.on_auth_failure = on_auth_failure
        auth_manager.on_device_paired = on_device_paired
        
        # Start auth manager
        await auth_manager.start()
        
        try:
            # Test PIN authentication flow
            device_id = "test_device_001"
            
            # Create challenge
            challenge = await auth_manager.create_auth_challenge(device_id, "123456")
            print(f"Challenge created: {challenge.challenge_id}")
            
            # Test wrong PIN
            result, session = await auth_manager.verify_auth_response(
                challenge.challenge_id, "654321"
            )
            print(f"Wrong PIN result: {result.name}")
            
            # Create new challenge
            challenge = await auth_manager.create_auth_challenge(device_id, "123456")
            
            # Test correct PIN
            result, session = await auth_manager.verify_auth_response(
                challenge.challenge_id, "123456"
            )
            print(f"Correct PIN result: {result.name}")
            
            if session:
                print(f"Session created: {session.session_id}")
                print(f"Permissions: {session.permissions}")
                
                # Test device pairing
                device_info = await auth_manager.pair_device(
                    device_id, "Test Device", "AA:BB:CC:DD:EE:FF", trust_level=2
                )
                print(f"Device paired with trust level: {device_info.trust_level}")
                
                # Check updated permissions
                updated_session = auth_manager.get_session(session.session_id)
                if updated_session:
                    print(f"Updated permissions: {updated_session.permissions}")
                
                # Test permission checking
                has_read = auth_manager.has_permission(session.session_id, "read")
                has_admin = auth_manager.has_permission(session.session_id, "admin")
                print(f"Has read permission: {has_read}")
                print(f"Has admin permission: {has_admin}")
            
            # Print statistics
            stats = auth_manager.get_statistics()
            print(f"Statistics: {stats}")
            
        finally:
            # Cleanup
            await auth_manager.cleanup()
    
    # Run test
    asyncio.run(test_auth_manager())